#!/usr/bin/env python3

import os
import re
import matplotlib.pyplot as plt
import numpy as np
from collections import defaultdict
import glob

def parse_log_file(filepath):
    """Parse a log file and extract performance data"""
    data = {
        'array_size': 0,
        'device': '',
        'without_streams': 0.0,
        'with_streams': 0.0,
        'with_overlapped_streams': 0.0,
        'speedup_streams': 0.0,
        'speedup_overlapped': 0.0
    }
    
    try:
        with open(filepath, 'r') as f:
            content = f.read()
            
        # Extract array size
        array_size_match = re.search(r'Array size: (\d+) elements', content)
        if array_size_match:
            data['array_size'] = int(array_size_match.group(1))
            
        # Extract device name
        device_match = re.search(r'Device: (.+)', content)
        if device_match:
            data['device'] = device_match.group(1).strip()
            
        # Extract average results
        without_streams_match = re.search(r'Without streams: ([\d.]+) ms', content)
        with_streams_match = re.search(r'With streams: ([\d.]+) ms', content)
        overlapped_streams_match = re.search(r'With overlapped streams: ([\d.]+) ms', content)
        
        if without_streams_match:
            data['without_streams'] = float(without_streams_match.group(1))
        if with_streams_match:
            data['with_streams'] = float(with_streams_match.group(1))
        if overlapped_streams_match:
            data['with_overlapped_streams'] = float(overlapped_streams_match.group(1))
            
        # Extract speedup information
        speedup_streams_match = re.search(r'Streams vs Sequential: ([\d.]+)x speedup', content)
        speedup_overlapped_match = re.search(r'Overlapped Streams vs Sequential: ([\d.]+)x speedup', content)
        
        if speedup_streams_match:
            data['speedup_streams'] = float(speedup_streams_match.group(1))
        if speedup_overlapped_match:
            data['speedup_overlapped'] = float(speedup_overlapped_match.group(1))
            
    except Exception as e:
        print(f"Error parsing {filepath}: {e}")
        
    return data

def collect_data():
    """Collect data from all log files"""
    hip_data = []
    cuda_data = []
    
    # Parse HIP log files
    hip_files = glob.glob('output/hip_test_*.log')
    for filepath in hip_files:
        data = parse_log_file(filepath)
        if data['array_size'] > 0:
            hip_data.append(data)
    
    # Parse CUDA log files
    cuda_files = glob.glob('output/cuda_test_*.log')
    for filepath in cuda_files:
        data = parse_log_file(filepath)
        if data['array_size'] > 0:
            cuda_data.append(data)
    
    # Sort by array size
    hip_data.sort(key=lambda x: x['array_size'])
    cuda_data.sort(key=lambda x: x['array_size'])
    
    return hip_data, cuda_data

def plot_performance_comparison(hip_data, cuda_data):
    """Create comprehensive performance comparison plots"""
    
    # Get HIP and CUDA version info for filename
    hip_version = ""
    cuda_version = ""
    
    if hip_data and hip_data[0]['device']:
        hip_device = hip_data[0]['device'].replace(' ', '_').replace('/', '_')
        hip_version = f"HIP_{hip_device}"
    
    if cuda_data and cuda_data[0]['device']:
        cuda_device = cuda_data[0]['device'].replace(' ', '_').replace('/', '_')
        cuda_version = f"CUDA_{cuda_device}"
    
    # Create filename suffix based on available data
    if hip_data and cuda_data:
        filename_suffix = f"{hip_version}_vs_{cuda_version}"
    elif hip_data:
        filename_suffix = hip_version
    elif cuda_data:
        filename_suffix = cuda_version
    else:
        filename_suffix = "comparison"
    
    # Create figure with subplots
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle('CUDA vs HIP Stream Performance Comparison', fontsize=16, fontweight='bold')
    
    # Extract data for plotting
    hip_sizes = [d['array_size'] for d in hip_data]
    cuda_sizes = [d['array_size'] for d in cuda_data]
    
    # Plot 1: Execution Time Comparison
    ax1.set_title('Execution Time Comparison (Lower is Better)', fontweight='bold')
    ax1.set_xlabel('Array Size (elements)')
    ax1.set_ylabel('Execution Time (ms)')
    ax1.set_xscale('log')
    ax1.set_yscale('log')
    ax1.grid(True, alpha=0.3)
    
    # Plot execution times
    if hip_data:
        hip_without = [d['without_streams'] for d in hip_data]
        hip_with = [d['with_streams'] for d in hip_data]
        hip_overlapped = [d['with_overlapped_streams'] for d in hip_data]
        
        ax1.plot(hip_sizes, hip_without, 'r-o', label='HIP Without Streams', linewidth=2)
        ax1.plot(hip_sizes, hip_with, 'r--s', label='HIP With Streams', linewidth=2)
        ax1.plot(hip_sizes, hip_overlapped, 'r:^', label='HIP Overlapped Streams', linewidth=2)
    
    if cuda_data:
        cuda_without = [d['without_streams'] for d in cuda_data]
        cuda_with = [d['with_streams'] for d in cuda_data]
        cuda_overlapped = [d['with_overlapped_streams'] for d in cuda_data]
        
        ax1.plot(cuda_sizes, cuda_without, 'b-o', label='CUDA Without Streams', linewidth=2)
        ax1.plot(cuda_sizes, cuda_with, 'b--s', label='CUDA With Streams', linewidth=2)
        ax1.plot(cuda_sizes, cuda_overlapped, 'b:^', label='CUDA Overlapped Streams', linewidth=2)
    
    ax1.legend()
    
    # Plot 2: Speedup Comparison
    ax2.set_title('Speedup Comparison (Higher is Better)', fontweight='bold')
    ax2.set_xlabel('Array Size (elements)')
    ax2.set_ylabel('Speedup (x)')
    ax2.set_xscale('log')
    ax2.grid(True, alpha=0.3)
    
    if hip_data:
        hip_speedup_streams = [d['speedup_streams'] for d in hip_data]
        hip_speedup_overlapped = [d['speedup_overlapped'] for d in hip_data]
        
        ax2.plot(hip_sizes, hip_speedup_streams, 'r--s', label='HIP Streams Speedup', linewidth=2)
        ax2.plot(hip_sizes, hip_speedup_overlapped, 'r:^', label='HIP Overlapped Speedup', linewidth=2)
    
    if cuda_data:
        cuda_speedup_streams = [d['speedup_streams'] for d in cuda_data]
        cuda_speedup_overlapped = [d['speedup_overlapped'] for d in cuda_data]
        
        ax2.plot(cuda_sizes, cuda_speedup_streams, 'b--s', label='CUDA Streams Speedup', linewidth=2)
        ax2.plot(cuda_sizes, cuda_speedup_overlapped, 'b:^', label='CUDA Overlapped Speedup', linewidth=2)
    
    ax2.axhline(y=1.0, color='k', linestyle='-', alpha=0.3, label='No Speedup')
    ax2.legend()
    
    # Plot 3: Direct Performance Comparison (Best Method)
    ax3.set_title('Best Performance Comparison', fontweight='bold')
    ax3.set_xlabel('Array Size (elements)')
    ax3.set_ylabel('Best Execution Time (ms)')
    ax3.set_xscale('log')
    ax3.set_yscale('log')
    ax3.grid(True, alpha=0.3)
    
    if hip_data:
        hip_best = [min(d['with_streams'], d['with_overlapped_streams']) for d in hip_data]
        ax3.plot(hip_sizes, hip_best, 'r-o', label='HIP Best Performance', linewidth=3, markersize=8)
    
    if cuda_data:
        cuda_best = [min(d['with_streams'], d['with_overlapped_streams']) for d in cuda_data]
        ax3.plot(cuda_sizes, cuda_best, 'b-o', label='CUDA Best Performance', linewidth=3, markersize=8)
    
    ax3.legend()
    
    # Plot 4: Efficiency Comparison (Performance per Element)
    ax4.set_title('Processing Efficiency (Time per Element)', fontweight='bold')
    ax4.set_xlabel('Array Size (elements)')
    ax4.set_ylabel('Time per Element (ms/element)')
    ax4.set_xscale('log')
    ax4.set_yscale('log')
    ax4.grid(True, alpha=0.3)
    
    if hip_data:
        hip_efficiency = [min(d['with_streams'], d['with_overlapped_streams']) / d['array_size'] * 1000 for d in hip_data]
        ax4.plot(hip_sizes, hip_efficiency, 'r-o', label='HIP Efficiency', linewidth=2, markersize=6)
    
    if cuda_data:
        cuda_efficiency = [min(d['with_streams'], d['with_overlapped_streams']) / d['array_size'] * 1000 for d in cuda_data]
        ax4.plot(cuda_sizes, cuda_efficiency, 'b-o', label='CUDA Efficiency', linewidth=2, markersize=6)
    
    ax4.legend()
    
    plt.tight_layout()
    
    # Save plot with version-specific filename
    plot_filename = f'output/performance_comparison_{filename_suffix}.png'
    plt.savefig(plot_filename, dpi=300, bbox_inches='tight')
    plt.show()
    
    return plot_filename

def generate_performance_report(hip_data, cuda_data):
    """Generate a detailed performance report"""
    
    # Get HIP and CUDA version info for filename
    hip_version = ""
    cuda_version = ""
    
    if hip_data and hip_data[0]['device']:
        hip_device = hip_data[0]['device'].replace(' ', '_').replace('/', '_')
        hip_version = f"HIP_{hip_device}"
    
    if cuda_data and cuda_data[0]['device']:
        cuda_device = cuda_data[0]['device'].replace(' ', '_').replace('/', '_')
        cuda_version = f"CUDA_{cuda_device}"
    
    # Create filename suffix based on available data
    if hip_data and cuda_data:
        filename_suffix = f"{hip_version}_vs_{cuda_version}"
    elif hip_data:
        filename_suffix = hip_version
    elif cuda_data:
        filename_suffix = cuda_version
    else:
        filename_suffix = "comparison"
    
    report = []
    report.append("=" * 80)
    report.append("CUDA vs HIP Stream Performance Analysis Report")
    report.append("=" * 80)
    report.append("")
    
    if hip_data:
        report.append("HIP Performance Summary:")
        report.append("-" * 40)
        for data in hip_data:
            report.append(f"Array Size: {data['array_size']:>10} elements")
            report.append(f"  Device: {data['device']}")
            report.append(f"  Without Streams:     {data['without_streams']:>8.2f} ms")
            report.append(f"  With Streams:        {data['with_streams']:>8.2f} ms")
            report.append(f"  With Overlapped:     {data['with_overlapped_streams']:>8.2f} ms")
            report.append(f"  Streams Speedup:     {data['speedup_streams']:>8.2f}x")
            report.append(f"  Overlapped Speedup:  {data['speedup_overlapped']:>8.2f}x")
            report.append("")
    
    if cuda_data:
        report.append("CUDA Performance Summary:")
        report.append("-" * 40)
        for data in cuda_data:
            report.append(f"Array Size: {data['array_size']:>10} elements")
            report.append(f"  Device: {data['device']}")
            report.append(f"  Without Streams:     {data['without_streams']:>8.2f} ms")
            report.append(f"  With Streams:        {data['with_streams']:>8.2f} ms")
            report.append(f"  With Overlapped:     {data['with_overlapped_streams']:>8.2f} ms")
            report.append(f"  Streams Speedup:     {data['speedup_streams']:>8.2f}x")
            report.append(f"  Overlapped Speedup:  {data['speedup_overlapped']:>8.2f}x")
            report.append("")
    
    # Performance comparison
    if hip_data and cuda_data:
        report.append("Direct Performance Comparison:")
        report.append("-" * 40)
        
        # Find common array sizes
        hip_sizes = {d['array_size'] for d in hip_data}
        cuda_sizes = {d['array_size'] for d in cuda_data}
        common_sizes = sorted(hip_sizes & cuda_sizes)
        
        for size in common_sizes:
            hip_item = next(d for d in hip_data if d['array_size'] == size)
            cuda_item = next(d for d in cuda_data if d['array_size'] == size)
            
            hip_best = min(hip_item['with_streams'], hip_item['with_overlapped_streams'])
            cuda_best = min(cuda_item['with_streams'], cuda_item['with_overlapped_streams'])
            
            ratio = hip_best / cuda_best if cuda_best > 0 else float('inf')
            
            report.append(f"Array Size: {size:>10} elements")
            report.append(f"  HIP Best:   {hip_best:>8.2f} ms")
            report.append(f"  CUDA Best:  {cuda_best:>8.2f} ms")
            report.append(f"  Ratio:      {ratio:>8.2f}x {'(HIP slower)' if ratio > 1 else '(CUDA slower)'}")
            report.append("")
    
    # Key insights
    report.append("Key Insights:")
    report.append("-" * 40)
    
    if hip_data:
        avg_hip_streams_speedup = np.mean([d['speedup_streams'] for d in hip_data])
        avg_hip_overlapped_speedup = np.mean([d['speedup_overlapped'] for d in hip_data])
        report.append(f"• HIP average streams speedup: {avg_hip_streams_speedup:.2f}x")
        report.append(f"• HIP average overlapped speedup: {avg_hip_overlapped_speedup:.2f}x")
    
    if cuda_data:
        avg_cuda_streams_speedup = np.mean([d['speedup_streams'] for d in cuda_data])
        avg_cuda_overlapped_speedup = np.mean([d['speedup_overlapped'] for d in cuda_data])
        report.append(f"• CUDA average streams speedup: {avg_cuda_streams_speedup:.2f}x")
        report.append(f"• CUDA average overlapped speedup: {avg_cuda_overlapped_speedup:.2f}x")
    
    report.append("")
    report.append("=" * 80)
    
    # Write report to file with version-specific filename
    report_filename = f'output/performance_report_{filename_suffix}.txt'
    with open(report_filename, 'w') as f:
        f.write('\n'.join(report))
    
    # Print to console
    print('\n'.join(report))
    
    return report_filename

def main():
    """Main function"""
    print("Analyzing CUDA vs HIP Stream Performance...")
    print("=" * 50)
    
    # Create output directory if it doesn't exist
    os.makedirs('output', exist_ok=True)
    
    # Collect performance data
    hip_data, cuda_data = collect_data()
    
    if not hip_data and not cuda_data:
        print("No performance data found. Please run the test scripts first:")
        print("  ./test_hip_sizes.sh")
        print("  ./test_cuda_sizes.sh")
        return
    
    print(f"Found {len(hip_data)} HIP test results")
    print(f"Found {len(cuda_data)} CUDA test results")
    print("")
    
    # Generate plots
    print("Generating performance plots...")
    plot_filename = plot_performance_comparison(hip_data, cuda_data)
    print(f"Performance plots saved to: {plot_filename}")
    
    # Generate report
    print("Generating performance report...")
    report_filename = generate_performance_report(hip_data, cuda_data)
    print(f"Performance report saved to: {report_filename}")
    
    print("\nAnalysis complete!")

if __name__ == "__main__":
    main()